import collections
import os
import platform
import time

import gymnasium
import numpy as np
from gymnasium.spaces import Box

import ogbench
from utils.datasets import Dataset


class EpisodeMonitor(gymnasium.Wrapper):
    """Environment wrapper to monitor episode statistics."""

    def __init__(self, env):
        super().__init__(env)
        self._reset_stats()
        self.total_timesteps = 0

    def _reset_stats(self):
        self.reward_sum = 0.0
        self.episode_length = 0
        self.start_time = time.time()

    def step(self, action):
        observation, reward, terminated, truncated, info = self.env.step(action)

        self.reward_sum += reward
        self.episode_length += 1
        self.total_timesteps += 1
        info['total'] = {'timesteps': self.total_timesteps}

        if terminated or truncated:
            info['episode'] = {}
            info['episode']['return'] = self.reward_sum
            info['episode']['length'] = self.episode_length
            info['episode']['duration'] = time.time() - self.start_time

        return observation, reward, terminated, truncated, info

    def reset(self, *args, **kwargs):
        self._reset_stats()
        return self.env.reset(*args, **kwargs)


class FrameStackWrapper(gymnasium.Wrapper):
    """Environment wrapper to stack observations."""

    def __init__(self, env, num_stack):
        super().__init__(env)

        self.num_stack = num_stack
        self.frames = collections.deque(maxlen=num_stack)

        low = np.concatenate([self.observation_space.low] * num_stack, axis=-1)
        high = np.concatenate([self.observation_space.high] * num_stack, axis=-1)
        self.observation_space = Box(low=low, high=high, dtype=self.observation_space.dtype)

    def get_observation(self):
        assert len(self.frames) == self.num_stack
        return np.concatenate(list(self.frames), axis=-1)

    def reset(self, **kwargs):
        ob, info = self.env.reset(**kwargs)
        for _ in range(self.num_stack):
            self.frames.append(ob)
        if 'goal' in info:
            info['goal'] = np.concatenate([info['goal']] * self.num_stack, axis=-1)
        return self.get_observation(), info

    def step(self, action):
        ob, reward, terminated, truncated, info = self.env.step(action)
        self.frames.append(ob)
        return self.get_observation(), reward, terminated, truncated, info


# utils/env_utils.py
# utils/env_utils.py
import os, glob
import ogbench
# 其余 import 保持不变

def _ensure_dataset_cached(dataset_name: str, dataset_dir: str):
    os.makedirs(dataset_dir, exist_ok=True)
    # 训练/验证数据通常命名为 <name>.npz 和 <name>-val.npz
    train_glob = glob.glob(os.path.join(dataset_dir, f"{dataset_name}.npz"))
    val_glob   = glob.glob(os.path.join(dataset_dir, f"{dataset_name}*-val.npz"))
    if train_glob and val_glob:
        print(f"[OGBench] cache found in {dataset_dir}, skip download.")
        return
    print(f"[OGBench] downloading {dataset_name} to {dataset_dir} ...")
    ogbench.download_datasets([dataset_name], dataset_dir=dataset_dir)

def make_env_and_datasets(dataset_name, frame_stack=None, dataset_dir=None):
    kwargs = {"compact_dataset": True}
    if dataset_dir is not None:
        _ensure_dataset_cached(dataset_name, dataset_dir)
        kwargs["dataset_dir"] = dataset_dir

    env, train_dict, val_dict = ogbench.make_env_and_datasets(dataset_name, **kwargs)

    # 不要再断言 train_dict['dataset_dir'] —— 返回字典里没有这个键
    train_dataset = Dataset.create(**train_dict)
    val_dataset   = Dataset.create(**val_dict) if val_dict is not None else None

    if frame_stack is not None:
        env = FrameStackWrapper(env, frame_stack)
    env.reset()
    return env, train_dataset, val_dataset
